import pathlib

import numpy as np
import quimb as qu
import quimb.tensor as qtn
from quimb.tensor import Tensor, TensorNetwork

I, X, Y, Z = np.eye(2), np.array([[0.0, 1], [1, 0]]), np.array([[0.0, -1j], [1j, 0]]), np.array([[1.0, 0], [0, -1]])

I_o1 = np.array([np.sqrt(2), 0, 0, 0])
I_o2 = np.identity(4)
Xp_o1 = np.array([1, 1, 0, 0]) / np.sqrt(2)
Xm_o1 = np.array([1, -1, 0, 0]) / np.sqrt(2)
Yp_o1 = np.array([1, 0, 1, 0]) / np.sqrt(2)
Ym_o1 = np.array([1, 0, -1, 0]) / np.sqrt(2)
Zp_o1 = np.array([1, 0, 0, 1]) / np.sqrt(2)
Zm_o1 = np.array([1, 0, 0, -1]) / np.sqrt(2)

meas_map = {
	'x+': Xp_o1, 'x-': Xm_o1,
	'y+': Yp_o1, 'y-': Ym_o1,
	'z+': Zp_o1, 'z-': Zm_o1,
}
"""Setting"""
cached = True
N_qubit, step = 10, 7
J, h, dt = 0.5236, 1.0, 0.05
max_bond, cutoff = 400, 1e-10
import json

with open('../splm_coeff_1.json', 'r') as f:
	data_dict1 = json.load(f)
	layer1_lam1_mat = np.asarray(data_dict1['lam1s'])
	layer1_lam2_mat = np.asarray(data_dict1['lam2s'])
with open('../splm_coeff_2.json', 'r') as f:
	data_dict2 = json.load(f)
	layer2_lam1_mat = np.asarray(data_dict2['lam1s'])
	layer2_lam2_mat = np.asarray(data_dict2['lam2s'])
# guess_layer1_lam1_mat = np.load('lam1_mat1 N_qubit=10.npy')
# guess_layer1_lam2_mat = np.load('lam2_mat1 N_qubit=10.npy')
# guess_layer2_lam1_mat = np.load('lam1_mat2 N_qubit=10.npy')
# guess_layer2_lam2_mat = np.load('lam2_mat2 N_qubit=10.npy')

guess_layer1_lam1_mat = np.load('lam1_mat1.npy')
guess_layer1_lam2_mat = np.load('lam2_mat1.npy')
guess_layer2_lam1_mat = np.load('lam1_mat2.npy')
guess_layer2_lam2_mat = np.load('lam2_mat2.npy')

# from QPT.utils import SPLM_coeff_plot
# import matplotlib.pyplot as plt
#
# _, axs = plt.subplots(2, 1, figsize=(14, 6.5), layout='constrained')
# SPLM_coeff_plot(N_qubit, layer1_lam1_mat, layer1_lam2_mat, ax=axs[0], title="real")
# SPLM_coeff_plot(N_qubit, guess_layer1_lam1_mat, guess_layer1_lam2_mat, ax=axs[1], title='guess')
# plt.show()
#
# _, axs = plt.subplots(2, 1, figsize=(14, 6.5), layout='constrained')
# SPLM_coeff_plot(N_qubit, layer2_lam1_mat, layer2_lam2_mat, ax=axs[0], title="real")
# SPLM_coeff_plot(N_qubit, guess_layer2_lam1_mat, guess_layer2_lam2_mat, ax=axs[1], title='guess')
# plt.show()
"""MPO representation of trotter step"""
site_tags = list(map("I{}".format, range(N_qubit)))

index_map1 = {old_upper_ind: new_upper_ind for old_upper_ind, new_upper_ind in
			  map(lambda _: ("k{}".format(_), "c{}".format(_)), range(N_qubit))}
index_map2 = {old_lower_ind: new_lower_ind for old_lower_ind, new_lower_ind in
			  map(lambda _: ("b{}".format(_), "c{}".format(_)), range(N_qubit))}

"""ideal trotter step"""
exist_step = step
while exist_step > 0:
	exist_trotter_step_filename = f"cache/trotter_step, step={exist_step}, J={J:.4f}, h={h:.2f}, dt={dt:.2f}, max_bond={max_bond}, cutoff={cutoff:.1e}"
	if pathlib.Path(exist_trotter_step_filename).exists():
		print(f"Existing File of trotter step at {exist_step} MPO Found")
		trotter_step = qu.load_from_disk(exist_trotter_step_filename)
		break
	exist_step -= 1
else:
	print(f"No Existing File of any trotter step MPO Found")
	trotter_step = TensorNetwork([Tensor(I_o2, inds=(f"k{i}", f"b{i}"), tags=[f"I{i}"]) for i in range(N_qubit)],
								 virtual=True)

if exist_step < step:
	from utils import trotter_step_MPO

	print("Loading one_trotter_step ...")
	one_trotter_step = trotter_step_MPO(N_qubit, J, h, dt, _max_bond=max_bond, _cutoff=cutoff)
	print("Loading Done.")

	for i in range(exist_step, step):
		print(f"Preparing trotter step at {i + 1} MPO ...")
		trotter_step = trotter_step.reindex(index_map1) | one_trotter_step.reindex(index_map2)
		qtn.tensor_network_1d_compress(trotter_step, max_bond=max_bond, cutoff=cutoff, site_tags=site_tags,
									   method='direct', inplace=True)
		if cached:
			trotter_step_filename = f"cache/trotter_step, step={i + 1}, J={J:.4f}, h={h:.2f}, dt={dt:.2f}, max_bond={max_bond}, cutoff={cutoff:.1e}"
			qu.save_to_disk(trotter_step, trotter_step_filename)
			print(f"Cache MPO of {i + 1} step to File '{trotter_step_filename}'.")
		print(f"MPO of step {i + 1} Done.")
	print(f"Finish.\n")

"""noisy trotter step test"""
exist_step = step
while exist_step > 0:
	exist_noisy_trotter_step_filename = f"cache/noisy_trotter_step, step={exist_step}, J={J:.4f}, h={h:.2f}, dt={dt:.2f}, max_bond={max_bond}, cutoff={cutoff:.1e}"
	if pathlib.Path(exist_noisy_trotter_step_filename).exists():
		print(f"Existing File of noisy trotter step at {exist_step} MPO Found")
		noisy_trotter_step = qu.load_from_disk(exist_noisy_trotter_step_filename)
		break
	exist_step -= 1
else:
	print(f"No Existing File of any noisy trotter step MPO Found")
	noisy_trotter_step = TensorNetwork([Tensor(I_o2, inds=(f"k{i}", f"b{i}"), tags=[f"I{i}"]) for i in range(N_qubit)],
									   virtual=True)

if exist_step < step:
	from utils import noisy_trotter_step_MPO
	print("Loading noisy_one_trotter_step ...")
	one_noisy_trotter_step = noisy_trotter_step_MPO(N_qubit, J, h, dt,
													layer1_lam1_mat, layer1_lam2_mat, layer2_lam1_mat, layer2_lam2_mat,
													_max_bond=max_bond, _cutoff=cutoff)
	print("Loading Done.")

	for i in range(exist_step, step):
		print(f"Preparing noisy trotter step at {i + 1} MPO ...")
		noisy_trotter_step = noisy_trotter_step.reindex(index_map1) | one_noisy_trotter_step.reindex(index_map2)
		qtn.tensor_network_1d_compress(noisy_trotter_step, max_bond=max_bond, cutoff=cutoff, site_tags=site_tags,
									   method='direct', inplace=True)
		if cached:
			noisy_trotter_step_filename = f"cache/noisy_trotter_step, step={i + 1}, J={J:.4f}, h={h:.2f}, dt={dt:.2f}, max_bond={max_bond}, cutoff={cutoff:.1e}"
			qu.save_to_disk(noisy_trotter_step, noisy_trotter_step_filename)
			print(f"Cache noisy MPO of {i + 1} step to File '{noisy_trotter_step_filename}'.")
		print(f"noisy MPO of step {i + 1} Done.")
	print(f"Finish.\n")

"""TEM trotter step test"""
exist_step = step
while exist_step > 0:
	exist_TEM_trotter_step_filename = f"cache/TEM_trotter_step, step={exist_step}, J={J:.4f}, h={h:.2f}, dt={dt:.2f}, max_bond={max_bond}, cutoff={cutoff:.1e}"
	if pathlib.Path(exist_TEM_trotter_step_filename).exists():
		print(f"Existing File of TEM trotter step at {exist_step} MPO Found")
		TEM_trotter_step = qu.load_from_disk(exist_TEM_trotter_step_filename)
		break
	exist_step -= 1
else:
	print(f"No Existing File of any TEM trotter step MPO Found")
	TEM_trotter_step = None

if exist_step < step:
	from utils import TEM_trotter_steps

	print("Loading TEM_one_trotter_step ...")
	TEM_one_trotter_step = TEM_trotter_steps(N_qubit, J, h, dt,
											 guess_layer1_lam1_mat, guess_layer1_lam2_mat,
											 guess_layer2_lam1_mat, guess_layer2_lam2_mat,
											 _max_bond=max_bond, _cutoff=cutoff)
	print("Loading Done.")

	for i in range(exist_step, step):
		print(f"Preparing TEM trotter step at {i + 1} MPO ...")
		TEM_trotter_step = TEM_one_trotter_step(TEM_trotter_step)
		qtn.tensor_network_1d_compress(TEM_trotter_step, max_bond=max_bond, cutoff=cutoff, site_tags=site_tags,
									   method='direct', inplace=True)
		if cached:
			TEM_trotter_step_filename = f"cache/TEM_trotter_step, step={i + 1}, J={J:.4f}, h={h:.2f}, dt={dt:.2f}, max_bond={max_bond}, cutoff={cutoff:.1e}"
			qu.save_to_disk(TEM_trotter_step, TEM_trotter_step_filename)
			print(f"Cache TEM MPO of {i + 1} step to File '{TEM_trotter_step_filename}'.")
		print(f"TEM MPO of step {i + 1} Done.")
	print(f"Finish.\n")

"""simple trotter step test"""
Zp_o1 = np.array([1, 0, 0, 1]) / np.sqrt(2)
Z_mea_o1 = np.array([0, 0, 0, np.sqrt(2)])
rho0_mps = TensorNetwork([Tensor(Zp_o1, inds=(f"b{i}",), tags=[f"I{i}"]) for i in range(N_qubit)], virtual=True)
Z_meas_mps = TensorNetwork([Tensor(Z_mea_o1, inds=(f"k{i}",), tags=[f"I{i}"]) for i in range(N_qubit)], virtual=True)
ideal_expectation_vals = [1.0]
noisy_expectation_vals = [1.0]
TEM_expectation_vals = [1.0]
for i in range(1, step + 1):
	trotter_step_filename = f"cache/trotter_step, step={i}, J={J:.4f}, h={h:.2f}, dt={dt:.2f}, max_bond={max_bond}, cutoff={cutoff:.1e}"
	trotter_step = qu.load_from_disk(trotter_step_filename)
	noisy_trotter_step_filename = f"cache/noisy_trotter_step, step={i}, J={J:.4f}, h={h:.2f}, dt={dt:.2f}, max_bond={max_bond}, cutoff={cutoff:.1e}"
	noisy_trotter_step = qu.load_from_disk(noisy_trotter_step_filename)
	TEM_trotter_step_filename = f"cache/TEM_trotter_step, step={i}, J={J:.4f}, h={h:.2f}, dt={dt:.2f}, max_bond={max_bond}, cutoff={cutoff:.1e}"
	TEM_trotter_step = qu.load_from_disk(TEM_trotter_step_filename)

	ideal_expectation_vals.append((rho0_mps | trotter_step | Z_meas_mps).contract())
	noisy_expectation_vals.append((rho0_mps | noisy_trotter_step | Z_meas_mps).contract())
	tem = noisy_trotter_step.reindex(index_map1) | TEM_trotter_step.reindex(index_map2)
	TEM_expectation_vals.append((rho0_mps | tem | Z_meas_mps).contract())

import matplotlib.pyplot as plt

plt.plot(np.arange(len(ideal_expectation_vals)), ideal_expectation_vals, 'k--', label='ideal')
plt.plot(np.arange(len(noisy_expectation_vals)), noisy_expectation_vals, color='orange', label='noisy')
plt.plot(np.arange(len(TEM_expectation_vals)), TEM_expectation_vals, color='green', label='TEM')
plt.xlim([0, step])
plt.ylim([-1.5, 1.5])
plt.xlabel('steps', fontsize=14)
plt.ylabel(r'$\langle Z^{\otimes 10}\rangle$', fontsize=14)
plt.legend()
plt.show()
